import numpy as np
from .bbox import bbox_iou

__all__ = ['corner2center', 'center2corner', 'gen_anchors', 'get_targets']

def corner2center(bbox, split=False):
    x1, y1, x2, y2 = bbox.T
    cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
    w, h = x2 - x1, y2 - y1
    if split:
        return cx, cy, w, h
    else:
        return np.stack((cx, cy, w, h), axis=1)

def center2corner(bbox, split=False):
    cx, cy, w, h = bbox.T
    x1, x2 = cx - w / 2, cx + w / 2
    y1, y2 = cy - h / 2, cy + h / 2
    if split:
        return x1, y1, x2, y2
    else:
        return np.stack((x1, y1, x2, y2), axis=1)

def gen_anchors(maps, strides, sizes, receptive_fields):
    anchors = []
    for mp, stride, size, repfile in zip(maps, strides, sizes, receptive_fields):
        height, width = mp.shape[2:]
        offset = (repfile - size) / 2
        offset_x = np.arange(width, dtype=np.float32)
        offset_y = np.arange(height, dtype=np.float32)
        offset_x, offset_y = np.meshgrid(offset_x, offset_y)
        anchor = np.stack((offset_x.ravel(), offset_y.ravel(),
                           offset_x.ravel(), offset_y.ravel()), axis=1)
        anchor = anchor * stride + offset
        anchor[:, 2:] += size
        anchors.append(anchor)
    return np.concatenate(anchors, axis=0)

def get_targets(anchors, label, stds=(0.1, 0.1, 0.2, 0.2)):
    targets = []
    ious = bbox_iou(anchors, label)
    gtbox = label[ious.argmax(1)]
    ious = ious.max(1)
    acx, acy, aw, ah = corner2center(anchors, True)
    gcx, gcy, gw, gh = corner2center(gtbox, True)
    tx = (gcx - acx) / aw / stds[0]
    ty = (gcy - acy) / ah / stds[1]
    tw = np.log(gw / aw) / stds[2]
    th = np.log(gh / ah) / stds[3]
    return np.stack((ious, tx, ty, tw, th), axis=1)